package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// 全局变量
var (
	serverURL string
	timeout   time.Duration
	verbose   bool
)

// 命令行参数
func init() {
	flag.StringVar(&serverURL, "server", "http://localhost:8080", "FastDNS 服务器地址")
	flag.DurationVar(&timeout, "timeout", 10*time.Second, "请求超时时间")
	flag.BoolVar(&verbose, "verbose", false, "显示详细输出")
}

// 主函数
func main() {
	// 解析命令行参数
	flag.Parse()

	// 检查是否有子命令
	args := flag.Args()
	if len(args) < 1 {
		printUsage()
		os.Exit(1)
	}

	// 根据子命令执行相应的操作
	switch args[0] {
	case "health":
		checkHealth()
	case "status":
		getStatus()
	case "upstream":
		handleUpstream(args[1:])
	case "cache":
		handleCache(args[1:])
	case "query":
		handleQuery(args[1:])
	case "benchmark":
		handleBenchmark(args[1:])
	case "help":
		printUsage()
	default:
		fmt.Printf("未知命令: %s\n", args[0])
		printUsage()
		os.Exit(1)
	}
}

// 打印使用帮助
func printUsage() {
	fmt.Println("FastDNS 命令行客户端")
	fmt.Println("\n用法:")
	fmt.Println("  fastdnsctl [选项] 命令 [参数]")
	fmt.Println("\n选项:")
	fmt.Println("  -server string    FastDNS 服务器地址 (默认 \"http://localhost:8080\")")
	fmt.Println("  -timeout duration 请求超时时间 (默认 10s)")
	fmt.Println("  -verbose          显示详细输出")
	fmt.Println("\n命令:")
	fmt.Println("  health            检查服务健康状态")
	fmt.Println("  status            获取服务运行状态")
	fmt.Println("  upstream list     列出所有上游 DNS 服务器")
	fmt.Println("  upstream add      添加上游 DNS 服务器")
	fmt.Println("  upstream remove   移除上游 DNS 服务器")
	fmt.Println("  cache status      获取缓存状态")
	fmt.Println("  cache clear       清除缓存")
	fmt.Println("  query             查询域名")
	fmt.Println("  benchmark         执行压力测试")
	fmt.Println("  help              显示帮助信息")
	fmt.Println("\n示例:")
	fmt.Println("  fastdnsctl health")
	fmt.Println("  fastdnsctl status")
	fmt.Println("  fastdnsctl upstream list")
	fmt.Println("  fastdnsctl upstream add --address 8.8.8.8:53 --protocol udp --timeout 5s --priority 1")
	fmt.Println("  fastdnsctl upstream remove --address 8.8.8.8:53")
	fmt.Println("  fastdnsctl cache status")
	fmt.Println("  fastdnsctl cache clear")
	fmt.Println("  fastdnsctl query --domain example.com --type A")
	fmt.Println("  fastdnsctl benchmark --domain example.com --type A --concurrent 10 --requests 100 --timeout 30s")
}

// 检查服务健康状态
func checkHealth() {
	url := fmt.Sprintf("%s/api/v1/health", serverURL)
	resp, err := sendRequest("GET", url, nil)
	if err != nil {
		fmt.Printf("检查健康状态失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	var result map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Printf("解析响应失败: %v\n", err)
		os.Exit(1)
	}

	status, ok := result["status"].(string)
	if !ok {
		fmt.Println("无效的响应格式")
		os.Exit(1)
	}

	if status == "healthy" {
		fmt.Println("服务状态: 健康")
	} else {
		fmt.Printf("服务状态: %s\n", status)
	}
}

// 获取服务运行状态
func getStatus() {
	url := fmt.Sprintf("%s/api/v1/status", serverURL)
	resp, err := sendRequest("GET", url, nil)
	if err != nil {
		fmt.Printf("获取服务状态失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	var result map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Printf("解析响应失败: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("服务状态:")
	fmt.Printf("  运行状态: %v\n", result["status"])
	fmt.Printf("  运行时间: %v\n", result["uptime"])
	fmt.Printf("  处理查询数: %v\n", result["queries_processed"])
	fmt.Printf("  缓存启用: %v\n", result["cache_enabled"])

	if result["cache_enabled"].(bool) {
		fmt.Printf("  缓存大小: %v\n", result["cache_size"])
		fmt.Printf("  缓存容量: %v\n", result["cache_capacity"])
		if hitRatio, ok := result["cache_hit_ratio"]; ok {
			fmt.Printf("  缓存命中率: %.2f%%\n", hitRatio.(float64)*100)
		}
	}
}

// 处理上游 DNS 服务器相关命令
func handleUpstream(args []string) {
	if len(args) < 1 {
		fmt.Println("缺少上游 DNS 服务器子命令")
		printUsage()
		os.Exit(1)
	}

	switch args[0] {
	case "list":
		listUpstreamServers()
	case "add":
		addUpstreamServer(args[1:])
	case "remove":
		removeUpstreamServer(args[1:])
	default:
		fmt.Printf("未知的上游 DNS 服务器子命令: %s\n", args[0])
		printUsage()
		os.Exit(1)
	}
}

// 列出所有上游 DNS 服务器
func listUpstreamServers() {
	url := fmt.Sprintf("%s/api/v1/dns/upstream", serverURL)
	resp, err := sendRequest("GET", url, nil)
	if err != nil {
		fmt.Printf("获取上游 DNS 服务器列表失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	var result map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Printf("解析响应失败: %v\n", err)
		os.Exit(1)
	}

	servers, ok := result["upstream_servers"].([]any)
	if !ok {
		fmt.Println("无效的响应格式")
		os.Exit(1)
	}

	if len(servers) == 0 {
		fmt.Println("没有配置上游 DNS 服务器")
		return
	}

	fmt.Println("上游 DNS 服务器列表:")
	for i, server := range servers {
		s := server.(map[string]any)
		fmt.Printf("  %d. 地址: %s\n", i+1, s["address"])
		fmt.Printf("     协议: %s\n", s["protocol"])
		fmt.Printf("     超时: %s\n", s["timeout"])
		fmt.Printf("     优先级: %d\n", int(s["priority"].(float64)))
		fmt.Printf("     状态: %s\n", s["status"])
		fmt.Printf("     延迟: %s\n", s["latency"])
		if i < len(servers)-1 {
			fmt.Println()
		}
	}
}

// 添加上游 DNS 服务器
func addUpstreamServer(args []string) {
	// 解析参数
	var address, protocol, timeout string
	var priority int

	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "--address":
			if i+1 < len(args) {
				address = args[i+1]
				i++
			}
		case "--protocol":
			if i+1 < len(args) {
				protocol = args[i+1]
				i++
			}
		case "--timeout":
			if i+1 < len(args) {
				timeout = args[i+1]
				i++
			}
		case "--priority":
			if i+1 < len(args) {
				_, err := fmt.Sscanf(args[i+1], "%d", &priority)
				if err != nil {
					fmt.Printf("无效的优先级: %s\n", args[i+1])
					os.Exit(1)
				}
				i++
			}
		}
	}

	// 验证参数
	if address == "" {
		fmt.Println("缺少必要参数: --address")
		os.Exit(1)
	}
	if protocol == "" {
		protocol = "udp" // 默认协议
	}
	if timeout == "" {
		timeout = "5s" // 默认超时
	}

	// 构建请求体
	requestBody := map[string]any{
		"address":  address,
		"protocol": protocol,
		"timeout":  timeout,
		"priority": priority,
	}

	// 发送请求
	url := fmt.Sprintf("%s/api/v1/dns/upstream", serverURL)
	resp, err := sendRequest("POST", url, requestBody)
	if err != nil {
		fmt.Printf("添加上游 DNS 服务器失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	// 处理响应
	var result map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Printf("解析响应失败: %v\n", err)
		os.Exit(1)
	}

	fmt.Println(result["message"])
}

// 移除上游 DNS 服务器
func removeUpstreamServer(args []string) {
	// 解析参数
	var address string

	for i := 0; i < len(args); i++ {
		if args[i] == "--address" && i+1 < len(args) {
			address = args[i+1]
			break
		}
	}

	// 验证参数
	if address == "" {
		fmt.Println("缺少必要参数: --address")
		os.Exit(1)
	}

	// 发送请求
	url := fmt.Sprintf("%s/api/v1/dns/upstream?address=%s", serverURL, address)
	resp, err := sendRequest("DELETE", url, nil)
	if err != nil {
		fmt.Printf("移除上游 DNS 服务器失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	// 处理响应
	var result map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Printf("解析响应失败: %v\n", err)
		os.Exit(1)
	}

	fmt.Println(result["message"])
}

// 处理缓存相关命令
func handleCache(args []string) {
	if len(args) < 1 {
		fmt.Println("缺少缓存子命令")
		printUsage()
		os.Exit(1)
	}

	switch args[0] {
	case "status":
		getCacheStatus()
	case "clear":
		clearCache()
	default:
		fmt.Printf("未知的缓存子命令: %s\n", args[0])
		printUsage()
		os.Exit(1)
	}
}

// 获取缓存状态
func getCacheStatus() {
	url := fmt.Sprintf("%s/api/v1/dns/cache", serverURL)
	resp, err := sendRequest("GET", url, nil)
	if err != nil {
		fmt.Printf("获取缓存状态失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	var result map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Printf("解析响应失败: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("缓存状态:")
	fmt.Printf("  当前大小: %v\n", result["size"])
	fmt.Printf("  最大容量: %v\n", result["capacity"])
	fmt.Printf("  TTL: %v\n", result["ttl"])
}

// 清除缓存
func clearCache() {
	url := fmt.Sprintf("%s/api/v1/dns/cache", serverURL)
	resp, err := sendRequest("DELETE", url, nil)
	if err != nil {
		fmt.Printf("清除缓存失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	var result map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Printf("解析响应失败: %v\n", err)
		os.Exit(1)
	}

	fmt.Println(result["message"])
}

// 处理 DNS 查询命令
func handleQuery(args []string) {
	// 解析参数
	var domain, recordType string

	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "--domain":
			if i+1 < len(args) {
				domain = args[i+1]
				i++
			}
		case "--type":
			if i+1 < len(args) {
				recordType = args[i+1]
				i++
			}
		}
	}

	// 验证参数
	if domain == "" {
		fmt.Println("缺少必要参数: --domain")
		os.Exit(1)
	}

	// 默认查询 A 记录
	if recordType == "" {
		recordType = "A"
	}

	// 验证记录类型
	validTypes := map[string]bool{
		"A":     true,
		"AAAA":  true,
		"CNAME": true,
		"MX":    true,
		"TXT":   true,
		"NS":    true,
		"SOA":   true,
		"SRV":   true,
		"PTR":   true,
		"CAA":   true,
	}

	if !validTypes[recordType] {
		fmt.Printf("不支持的记录类型: %s\n", recordType)
		fmt.Println("支持的记录类型: A, AAAA, CNAME, MX, TXT, NS, SOA, SRV, PTR, CAA")
		os.Exit(1)
	}

	// 构建请求体
	requestBody := map[string]string{
		"domain": domain,
		"type":   recordType,
	}

	// 发送请求
	url := fmt.Sprintf("%s/api/v1/dns/query", serverURL)
	resp, err := sendRequest("POST", url, requestBody)
	if err != nil {
		fmt.Printf("DNS 查询失败: %v\n", err)
		os.Exit(1)
	}
	defer resp.Body.Close()

	// 处理响应
	var result map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		fmt.Printf("解析响应失败: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("域名: %s\n", result["domain"])
	fmt.Printf("记录类型: %s\n", result["type"])

	records, ok := result["records"]
	if !ok {
		fmt.Println("无效的响应格式")
		os.Exit(1)
	}

	// 根据记录类型处理不同的响应格式
	switch recordType {
	case "A", "AAAA":
		// 处理 IP 地址记录
		ipList, ok := records.([]any)
		if !ok {
			fmt.Println("无效的响应格式")
			os.Exit(1)
		}

		fmt.Println("IP 地址:")
		for _, ip := range ipList {
			fmt.Printf("  %s\n", ip)
		}

	case "MX":
		// 处理 MX 记录
		mxList, ok := records.([]any)
		if !ok {
			fmt.Println("无效的响应格式")
			os.Exit(1)
		}

		fmt.Println("MX 记录:")
		for _, mx := range mxList {
			mxRecord := mx.(map[string]any)
			fmt.Printf("  主机: %s, 优先级: %d\n",
				mxRecord["host"], int(mxRecord["priority"].(float64)))
		}

	default:
		// 处理其他类型的记录
		recordList, ok := records.([]any)
		if !ok {
			fmt.Println("无效的响应格式")
			os.Exit(1)
		}

		fmt.Printf("%s 记录:\n", recordType)
		for _, record := range recordList {
			recordMap, ok := record.(map[string]any)
			if !ok {
				continue
			}

			fmt.Printf("  名称: %s, 值: %s\n",
				recordMap["name"], recordMap["value"])
		}
	}

	// 显示 TTL
	if ttl, ok := result["ttl"].(float64); ok {
		fmt.Printf("TTL: %.0f 秒\n", ttl)
	}
}

// 发送 HTTP 请求
func sendRequest(method, url string, body any) (*http.Response, error) {
	var reqBody io.Reader
	if body != nil {
		jsonData, err := json.Marshal(body)
		if err != nil {
			return nil, fmt.Errorf("序列化请求体失败: %v", err)
		}
		reqBody = bytes.NewBuffer(jsonData)
	}

	// 创建请求
	req, err := http.NewRequest(method, url, reqBody)
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %v", err)
	}

	// 设置请求头
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	req.Header.Set("Accept", "application/json")

	// 创建 HTTP 客户端
	client := &http.Client{
		Timeout: timeout,
	}

	// 打印详细信息
	if verbose {
		fmt.Printf("发送请求: %s %s\n", method, url)
		if body != nil {
			jsonData, _ := json.MarshalIndent(body, "", "  ")
			fmt.Printf("请求体:\n%s\n", string(jsonData))
		}
	}

	// 发送请求
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("发送请求失败: %v", err)
	}

	// 检查响应状态码
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		defer resp.Body.Close()

		// 尝试读取错误消息
		bodyBytes, err := io.ReadAll(resp.Body)
		if err != nil {
			return nil, fmt.Errorf("请求失败: HTTP %d", resp.StatusCode)
		}

		return nil, fmt.Errorf("请求失败: HTTP %d - %s", resp.StatusCode, strings.TrimSpace(string(bodyBytes)))
	}

	// 打印详细信息
	if verbose {
		fmt.Printf("收到响应: HTTP %d\n", resp.StatusCode)
	}

	return resp, nil
}

// 处理压力测试命令
func handleBenchmark(args []string) {
	// 解析参数
	var domain, recordType string
	concurrent := 10                 // 默认并发数
	requests := 100                  // 默认请求数
	benchTimeout := 30 * time.Second // 默认测试超时时间

	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "--domain":
			if i+1 < len(args) {
				domain = args[i+1]
				i++
			}
		case "--type":
			if i+1 < len(args) {
				recordType = args[i+1]
				i++
			}
		case "--concurrent", "-c":
			if i+1 < len(args) {
				_, err := fmt.Sscanf(args[i+1], "%d", &concurrent)
				if err != nil || concurrent < 1 {
					fmt.Printf("无效的并发数: %s\n", args[i+1])
					os.Exit(1)
				}
				i++
			}
		case "--requests", "-n":
			if i+1 < len(args) {
				_, err := fmt.Sscanf(args[i+1], "%d", &requests)
				if err != nil || requests < 1 {
					fmt.Printf("无效的请求数: %s\n", args[i+1])
					os.Exit(1)
				}
				i++
			}
		case "--timeout":
			if i+1 < len(args) {
				var err error
				benchTimeout, err = time.ParseDuration(args[i+1])
				if err != nil {
					fmt.Printf("无效的超时时间: %s\n", args[i+1])
					os.Exit(1)
				}
				i++
			}
		}
	}

	// 验证参数
	if domain == "" {
		fmt.Println("缺少必要参数: --domain")
		os.Exit(1)
	}

	// 默认查询 A 记录
	if recordType == "" {
		recordType = "A"
	}

	// 验证记录类型
	validTypes := map[string]bool{
		"A":     true,
		"AAAA":  true,
		"CNAME": true,
		"MX":    true,
		"TXT":   true,
		"NS":    true,
		"SOA":   true,
		"SRV":   true,
		"PTR":   true,
		"CAA":   true,
	}

	if !validTypes[recordType] {
		fmt.Printf("不支持的记录类型: %s\n", recordType)
		fmt.Println("支持的记录类型: A, AAAA, CNAME, MX, TXT, NS, SOA, SRV, PTR, CAA")
		os.Exit(1)
	}

	// 构建请求体
	requestBody := map[string]string{
		"domain": domain,
		"type":   recordType,
	}

	// 创建结果通道
	resultCh := make(chan *benchmarkResult, requests)

	// 创建统计数据
	var totalRequests, successRequests, failedRequests int64
	var totalLatency int64
	latencies := make([]time.Duration, 0, requests)
	var latenciesMutex sync.Mutex

	// 创建并发控制
	wg := sync.WaitGroup{}
	semaphore := make(chan struct{}, concurrent)

	// 记录开始时间
	startTime := time.Now()

	// 创建上下文和取消函数
	ctx, cancel := context.WithTimeout(context.Background(), benchTimeout)
	defer cancel()

	fmt.Printf("开始压力测试: 域名=%s, 类型=%s, 并发=%d, 请求数=%d\n",
		domain, recordType, concurrent, requests)

	// 启动工作协程
	for i := 0; i < requests; i++ {
		select {
		case <-ctx.Done():
			// 超时或取消
			goto DONE
		case semaphore <- struct{}{}:
			// 获取了信号量，可以执行
		}

		wg.Add(1)
		go func(reqNum int) {
			defer wg.Done()
			defer func() { <-semaphore }()

			// 发送请求并记录时间
			start := time.Now()
			url := fmt.Sprintf("%s/api/v1/dns/query", serverURL)
			resp, err := sendRequest("POST", url, requestBody)

			// 创建结果
			result := &benchmarkResult{
				reqNum:  reqNum,
				latency: time.Since(start),
			}

			// 记录请求结果
			atomic.AddInt64(&totalRequests, 1)

			if err != nil {
				result.err = err
				atomic.AddInt64(&failedRequests, 1)
			} else {
				defer resp.Body.Close()
				atomic.AddInt64(&successRequests, 1)
				atomic.AddInt64(&totalLatency, int64(result.latency))

				// 记录响应时间
				latenciesMutex.Lock()
				latencies = append(latencies, result.latency)
				latenciesMutex.Unlock()
			}

			// 发送结果
			resultCh <- result
		}(i + 1)
	}

	// 等待所有请求完成或超时
	go func() {
		wg.Wait()
		close(resultCh)
	}()

DONE:
	// 收集结果并显示进度
	var lastPercent int
	for result := range resultCh {
		// 计算完成百分比
		completed := atomic.LoadInt64(&totalRequests)
		percent := int(float64(completed) / float64(requests) * 100)

		// 每完成 10% 显示一次进度
		if percent/10 > lastPercent/10 {
			fmt.Printf("进度: %d%% (%d/%d)\n", percent, completed, requests)
			lastPercent = percent
		}

		// 详细输出模式下显示每个请求的结果
		if verbose {
			if result.err != nil {
				fmt.Printf("请求 #%d 失败: %v (耗时: %v)\n", result.reqNum, result.err, result.latency)
			} else {
				fmt.Printf("请求 #%d 成功 (耗时: %v)\n", result.reqNum, result.latency)
			}
		}
	}

	// 计算测试统计信息
	testDuration := time.Since(startTime)
	totalReq := atomic.LoadInt64(&totalRequests)
	successReq := atomic.LoadInt64(&successRequests)
	failedReq := atomic.LoadInt64(&failedRequests)

	// 排序响应时间以计算百分位数
	latenciesMutex.Lock()
	sort.Slice(latencies, func(i, j int) bool {
		return latencies[i] < latencies[j]
	})

	// 计算各个百分位数
	var p50, p90, p95, p99 time.Duration
	if len(latencies) > 0 {
		p50 = latencies[len(latencies)*50/100]
		p90 = latencies[len(latencies)*90/100]
		p95 = latencies[len(latencies)*95/100]
		p99 = latencies[len(latencies)*99/100]
	}
	latenciesMutex.Unlock()

	// 计算平均响应时间
	var avgLatency time.Duration
	if successReq > 0 {
		avgLatency = time.Duration(atomic.LoadInt64(&totalLatency) / successReq)
	}

	// 计算 QPS
	var qps float64
	if testDuration > 0 {
		qps = float64(successReq) / testDuration.Seconds()
	}

	// 输出测试结果
	fmt.Println("\n压力测试结果:")
	fmt.Printf("  域名: %s\n", domain)
	fmt.Printf("  记录类型: %s\n", recordType)
	fmt.Printf("  并发数: %d\n", concurrent)
	fmt.Printf("  总请求数: %d\n", totalReq)
	fmt.Printf("  成功请求数: %d\n", successReq)
	fmt.Printf("  失败请求数: %d\n", failedReq)
	fmt.Printf("  成功率: %.2f%%\n", float64(successReq)/float64(totalReq)*100)
	fmt.Printf("  总测试时间: %v\n", testDuration.Round(time.Millisecond))
	fmt.Printf("  平均响应时间: %v\n", avgLatency.Round(time.Microsecond))
	fmt.Printf("  响应时间百分位数:\n")
	fmt.Printf("    50%% 在 %v 内\n", p50.Round(time.Microsecond))
	fmt.Printf("    90%% 在 %v 内\n", p90.Round(time.Microsecond))
	fmt.Printf("    95%% 在 %v 内\n", p95.Round(time.Microsecond))
	fmt.Printf("    99%% 在 %v 内\n", p99.Round(time.Microsecond))
	fmt.Printf("  QPS: %.2f 请求/秒\n", qps)
}

// 压力测试结果
type benchmarkResult struct {
	reqNum  int
	latency time.Duration
	err     error
}
